package com.redis.limit;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.Pipeline;
import redis.clients.jedis.Response;

/**
 * Distributed Rate Limiting Strategies
 */
public class DistributeRateLimiter {

    private static final String PREFIX_KEY = "hist:";

    private final String action;

    private final String rowkey;

    private final int period;

    private final int maxCount;

    private final Jedis jedis;

    public DistributeRateLimiter(String action, int period, int maxCount, String host, int port) {
        this.action = action;
        this.rowkey = PREFIX_KEY + action + ":%s";
        this.period = period;
        this.maxCount = maxCount;
        this.jedis = new Jedis(host, port);
    }

    public boolean isActionAllowed(String id) {
        String key = String.format(rowkey, id);
        long nowTimestamp = System.currentTimeMillis() * 1000 + System.nanoTime() % 1000;
        Pipeline pipeline = jedis.pipelined();
        pipeline.zremrangeByScore(key, 0, nowTimestamp - 1000000L * period);
        pipeline.zadd(key, nowTimestamp, String.valueOf(nowTimestamp));
        Response<Long> response = pipeline.zcard(key);
        pipeline.expire(key, period + 1);
        pipeline.sync();
        Long count = response.get();
        pipeline.close();
        return count <= maxCount;
    }

    public void close() {
        jedis.close();
    }
}
